from flask import Flask, request, render_template
from werkzeug.utils import secure_filename
import os
import cv2
import numpy as np
import onnxruntime as ort
import yaml
import time


app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'static/'


@app.route("/")
def hello_world():
    return "<p>hello, world</p>"


@app.route('/upload')
def upload_file():
    return render_template('upload.html')


@app.route("/crowd_count", methods=['POST'])
def crowd_count():
    if request.method == 'POST':
        f = request.files['file']
        print(request.files)
        d_time = time.strftime('%Y%m%d%H%M%S', time.localtime(time.time()))
        path_ = 'static/'+d_time
        if not os.path.exists(path_):
            os.makedirs(path_)
        f.save(os.path.join(path_, secure_filename(f.filename)))
        configs = yaml.load(open('config.yaml').read(), Loader=yaml.FullLoader)
        model = ort.InferenceSession(configs['model'])
        img = preprocess(cv2.imread(os.path.join(path_, secure_filename(f.filename))), configs['input_size'])
        outputs_onnx = model.run(None, {'input': img})[0]
        pred = np.sum(outputs_onnx / 100)
        nums = int(pred + 0.5)
        level = 'comfortable'
        for key, value in configs['level'].items():
            if len(value) == 1:
                if nums >= value[0]:
                    level = key
                else:
                    continue
            elif len(value) == 2:
                if nums >= value[0] and nums <= value[1]:
                    level = key
                else:
                    continue
            else:
                raise Exception("Invalid value!", value)
        text = f'nums: {nums}  level: {level}'
        return text
    else:
        return "METHOD NOT ALLOWED"


@app.route("/res", methods=['POST'])
def res():
    if request.method == 'POST':
        f = request.files['file']
        d_time = time.strftime('%Y%m%d%H%M%S', time.localtime(time.time()))
        path_ = 'static/'+d_time
        if not os.path.exists(path_):
            os.makedirs(path_)
        f.save(os.path.join(path_, secure_filename(f.filename)))
        configs = yaml.load(open('config.yaml').read(), Loader=yaml.FullLoader)
        model = ort.InferenceSession(configs['model'])
        img = preprocess(cv2.imread(os.path.join(path_, secure_filename(f.filename))), configs['input_size'])
        outputs_onnx = model.run(None, {'input': img})[0]
        pred = np.sum(outputs_onnx / 100)
        nums = int(pred + 0.5)
        level = 'comfortable'
        for key, value in configs['level'].items():
            if len(value) == 1:
                if nums >= value[0]:
                    level = key
                else:
                    continue
            elif len(value) == 2:
                if nums >= value[0] and nums <= value[1]:
                    level = key
                else:
                    continue
            else:
                raise Exception("Invalid value!", value)
        text = f'nums: {nums}  level: {level}'
        img_ori = cv2.putText(cv2.imread(os.path.join(path_, secure_filename(f.filename))), text, (50, 50), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255), 1)
        save_file_path = os.path.join(path_, secure_filename(f.filename).replace(secure_filename(f.filename).split('.')[-1],'-success.'+secure_filename(f.filename).split('.')[-1]))
        cv2.imwrite(save_file_path, img_ori)
        return render_template('res.html', img_src=save_file_path)
    else:
        return render_template('upload.html')


def preprocess(img, size):
    img_w, img_h = img.shape[1], img.shape[0]
    w, h = size
    new_w = int(img_w * min(w * 1. / img_w, h * 1. / img_h))
    new_h = int(img_h * min(w * 1. / img_w, h * 1. / img_h))
    resized_image = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_CUBIC)
    canvas = np.full((h, w, 3), 0)
    canvas[(h - new_h) // 2:(h - new_h) // 2 + new_h, (w - new_w) // 2:(w - new_w) // 2 + new_w, :] = resized_image
    input_array = canvas.transpose([2, 0, 1]) / 255.
    return input_array[np.newaxis,:].astype(np.float32)

def create_template(host,port):
    html_template = '''<!DOCTYPE html>
                    <html lang="en">
                    <head>
                        <meta charset="UTF-8">
                        <title>pic upload</title>
                    </head>
                    <body>
                        <form action="{0}/res" method="POST" enctype="multipart/form-data">
                            <input type="file" name="file" />
                            <input type="submit" name="submit">
                        </form>
                    </body>
                    </html>'''
    html_page = html_template.format('{0}:{1}'.format(host,port))
    with open('templates/upload.html','w') as f:
        f.write(html_page)

if __name__ == '__main__':
    configs = yaml.load(open('config.yaml').read(), Loader=yaml.FullLoader)
    host = '127.0.0.1'
    port = 8080
    try:
        server_configs = configs['server']
        host = server_configs['ip']
        port = server_configs['port']
    except Exception:
        pass
    # create_template(host,port)
    app.debug = True
    app.run(host=host, port=port)
    app.run(debug=True)